package com.deepseekrag.config;

import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

/**
 * 一般不是多供应商模型的话，使用自动配置就可以了，无需手动配置
 * 自动配置用spring-ai-starter-model-ollama
 * 手动配置用spring-ai-ollama
 */
@Configuration
public class RAGEmbeddingConfig {

    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter(512,400,5,800,true);
    }

    @Bean
    public OllamaEmbeddingModel ollamaEmbeddingModel() {
        OllamaApi ollamaApi = OllamaApi.builder().baseUrl("http://localhost:11434").build();

        OllamaOptions options = OllamaOptions.builder()
                .model(OllamaModel.NOMIC_EMBED_TEXT.id())
                .build();
        return OllamaEmbeddingModel.builder().ollamaApi(ollamaApi)
                .defaultOptions(options).build();
    }

    @Bean
    public PgVectorStore vectorStore(JdbcTemplate jdbcTemplate, OllamaEmbeddingModel ollamaEmbeddingModel) {
        return PgVectorStore.builder(jdbcTemplate,ollamaEmbeddingModel).build();
    }

    @Bean
    public SimpleVectorStore simpleVectorStore(OllamaEmbeddingModel ollamaEmbeddingModel) {
        return SimpleVectorStore.builder(ollamaEmbeddingModel).build();
    }
}
